import requests
import re
import argparse
from packaging.version import parse, Version
from typing import Dict, Tuple, List, Optional
from concurrent.futures import ThreadPoolExecutor, as_completed

PYPI_API_URL = "https://pypi.org/pypi/{package}/json"
USER_AGENT = "python-dependency-updater/1.0"

def parse_requirements(file_path: str) -> Dict[str, str]:
    """解析requirements.txt文件，返回包名和版本的字典"""
    requirements = {}
    pattern = re.compile(r'^(?P<name>[\w\-\.]+)(?P<spec>.*)$')
    
    with open(file_path, 'r') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#'):
                continue
            
            # 处理包含额外索引的复杂行
            if line.startswith('-e') or line.startswith('--') or '@' in line:
                requirements[line] = None  # 保留原始行
                continue
            
            match = pattern.match(line)
            if match:
                name = match.group('name').lower()
                spec = match.group('spec').strip()
                requirements[name] = spec if spec else None
            else:
                requirements[line] = None  # 保留无法解析的行
    
    return requirements

def get_latest_version(package: str) -> Tuple[Optional[Version], Optional[str]]:
    """通过PyPI API获取包的最新稳定版本"""
    try:
        response = requests.get(
            PYPI_API_URL.format(package=package),
            headers={'User-Agent': USER_AGENT},
            timeout=5
        )
        response.raise_for_status()
        data = response.json()
        
        # 过滤预发布版本
        releases = [parse(v) for v in data['releases'] 
                   if not parse(v).is_prerelease]
        if not releases:
            return None, None
            
        latest = max(releases)
        return latest, str(latest)
    except Exception as e:
        print(f"警告: 获取 {package} 信息失败 - {str(e)}")
        return None, None

def check_updates(requirements: Dict[str, str]) -> List[Tuple[str, str, str, str]]:
    """并发检查所有依赖的更新"""
    updates = []
    with ThreadPoolExecutor(max_workers=10) as executor:
        future_map = {}
        for pkg, current_spec in requirements.items():
            # 跳过复杂行和无法解析的行
            if not current_spec or not re.match(r'^[\w\-\.]+$', pkg):
                continue
            future_map[executor.submit(get_latest_version, pkg)] = (pkg, current_spec)
        
        for future in as_completed(future_map):
            pkg, current_spec = future_map[future]
            try:
                latest_version, latest_str = future.result()
                if not latest_version:
                    continue
                
                # 提取当前版本号 (忽略版本操作符)
                current_version_str = current_spec.replace('=', '').replace('~', '').strip()
                try:
                    current_version = parse(current_version_str)
                except:
                    current_version = None
                
                if current_version and latest_version > current_version:
                    updates.append((pkg, current_spec, f"=={latest_str}", str(latest_version)))
            except Exception as e:
                print(f"处理 {pkg} 时出错: {str(e)}")
    
    return updates

def update_requirements_file(input_file: str, output_file: str, updates: Dict[str, str]):
    """生成更新后的requirements文件"""
    with open(input_file, 'r') as fin, open(output_file, 'w') as fout:
        for line in fin:
            line_stripped = line.strip()
            if not line_stripped or line_stripped.startswith('#'):
                fout.write(line)
                continue
            
            # 尝试匹配包名
            pkg_match = re.match(r'^([\w\-\.]+)', line_stripped)
            if pkg_match:
                pkg_name = pkg_match.group(1).lower()
                if pkg_name in updates:
                    # 保留原始行的注释和格式
                    new_line = re.sub(
                        r'([=~]=?)[\w\.\*]+', 
                        f'\\g<1>{updates[pkg_name].split("==")[-1]}', 
                        line,
                        1
                    )
                    fout.write(new_line)
                    continue
            
            fout.write(line)

def main():
    parser = argparse.ArgumentParser(description='Python依赖更新工具')
    parser.add_argument('--file', default='requirements.txt', help='输入requirements文件')
    parser.add_argument('--output', default='requirements_new.txt', help='输出文件')
    parser.add_argument('--interactive', action='store_true', help='交互模式')
    parser.add_argument('--level', choices=['major', 'minor', 'patch'], 
                        default='minor', help='允许更新的版本级别')
    args = parser.parse_args()
    
    print(f"解析 {args.file}...")
    requirements = parse_requirements(args.file)
    print(f"发现 {len(requirements)} 个依赖项")
    
    print("检查更新...")
    available_updates = check_updates(requirements)
    if not available_updates:
        print("所有依赖都是最新的!")
        return
    
    print("\n发现可用更新:")
    updates_to_apply = {}
    for i, (pkg, current, new_spec, new_ver) in enumerate(available_updates, 1):
        print(f"{i}. {pkg}: {current} → {new_ver}")
    
    if args.interactive:
        print("\n选择更新选项:")
        print("1. 全部更新")
        print("2. 逐个确认")
        print("3. 跳过所有")
        choice = input("输入选择 (1/2/3): ").strip()
        
        if choice == '3':
            print("未进行任何更新")
            return
            
        for pkg, current, new_spec, new_ver in available_updates:
            if choice == '1' or (choice == '2' and input(f"更新 {pkg} ({current} → {new_ver})? [y/N] ").lower() == 'y'):
                updates_to_apply[pkg] = new_spec
    else:
        # 自动模式根据更新级别过滤
        for pkg, current, new_spec, new_ver in available_updates:
            current_ver = parse(current.replace('=', '').replace('~', '').strip())
            latest_ver = parse(new_ver)
            
            # 版本号分段比较
            if args.level == 'patch' and current_ver.release[:2] != latest_ver.release[:2]:
                continue  # 次要版本不同时跳过
            elif args.level == 'minor' and current_ver.release[0] != latest_ver.release[0]:
                continue  # 主版本不同时跳过
                
            updates_to_apply[pkg] = new_spec
    
    if updates_to_apply:
        update_requirements_file(args.file, args.output, updates_to_apply)
        print(f"\n更新完成! 已写入 {args.output}")
        print(f"更新了 {len(updates_to_apply)}/{len(available_updates)} 个依赖")
    else:
        print("没有符合更新策略的依赖")

if __name__ == "__main__":
    main()